package org.apache.zookeeper.server;

import org.apache.jute.Record;
import org.apache.jute.OutputArchive;
import org.apache.jute.InputArchive;
import org.apache.jute.BinaryOutputArchive;
import org.apache.zookeeper.proto.ReplyHeader;
import org.apache.zookeeper.proto.PrepareDataRequest;
import org.apache.zookeeper.server.paxos2.Configuration;
import org.apache.zookeeper.ZooDefs;

import java.nio.ByteBuffer;
import java.util.HashMap;
import java.util.TreeMap;
import java.io.IOException;
import java.io.ByteArrayOutputStream;

/**
 * Created by IntelliJ IDEA.
 * User: wangw
 * Date: 2011-6-5
 * Time: 10:29:41
 * To change this template use File | Settings | File Templates.
 */
public class ReplyCache {

    private static HashMap<Long, TreeMap<Integer, ByteBuffer>> cache = new HashMap<Long, TreeMap<Integer, ByteBuffer>>();

    private final static byte fourBytes[] = new byte[4];

    static private ByteBuffer serializeReply(ReplyHeader h, Record r, String tag) {
        try {
            ByteArrayOutputStream baos = new ByteArrayOutputStream();
            // Make space for length
            BinaryOutputArchive bos = BinaryOutputArchive.getArchive(baos);
            try {
                baos.write(fourBytes);
                bos.writeRecord(h, "header");
                if (r != null) {
                    bos.writeRecord(r, tag);
                }
                baos.close();
            } catch (IOException e) {
                e.printStackTrace();
            }
            byte b[] = baos.toByteArray();
            ByteBuffer bb = ByteBuffer.wrap(b);
            bb.putInt(b.length - 4).rewind();
            return bb;
        } catch (Exception e) {
            e.printStackTrace();
            return null;
        }
    }

    public synchronized static void addReply(long sessionId, int cxid, ReplyHeader h, Record r, String tag) {
        ByteBuffer buffer = serializeReply(h, r, tag);
        if (!cache.containsKey(sessionId))
            cache.put(sessionId, new TreeMap<Integer, ByteBuffer>());
        TreeMap<Integer, ByteBuffer> tmp = cache.get(sessionId);
        if (tmp.containsKey(cxid)){
            System.out.println("Reply cache already contains " + sessionId + " " + cxid);
	    return;
	}
	if(Configuration.debug)
	    System.out.println("put "+sessionId+" "+cxid+" into reply cache");
        tmp.put(cxid, buffer);
        if (tmp.size() > Configuration.clientOutStanding)
            tmp.remove(tmp.firstKey());
    }

    public synchronized static ByteBuffer getReply(long sessionId, int cxid) {
        if (!cache.containsKey(sessionId))
            return null;
        return cache.get(sessionId).get(cxid);
    }

    public synchronized static int getLastCxid(long sessionId) {
        if (!cache.containsKey(sessionId))
            return -1;
        else
            return cache.get(sessionId).lastKey();
    }

    public synchronized static int getFirstCxid(long sessionId) {
        if (!cache.containsKey(sessionId))
            return -1;
        else
            return cache.get(sessionId).firstKey();
    }


    public synchronized static void serialize(OutputArchive oa, String tag) throws IOException {
        //synchronized (cache) {
            oa.writeInt(cache.size(), "size");
            for (Long sessionId : cache.keySet()) {
                oa.writeLong(sessionId, "sessionId");
                TreeMap<Integer, ByteBuffer> client = cache.get(sessionId);
                oa.writeInt(client.size(), "clientsize");
                for (Integer xid : client.keySet()) {
                    oa.writeInt(xid, "xid");
                    oa.writeBuffer(client.get(xid).array(), "data");
		    if(Configuration.debug)
                        System.out.println("serialize replyCache" + sessionId + " " + xid);
                }
            }
        //}
    }

    public synchronized static void deserialize(InputArchive ia, String tag) throws IOException {
        //synchronized (cache) {
            cache.clear();
            int size = ia.readInt("size");
            System.out.println("deserialize reply cache size =" + size);
            for (int i = 0; i < size; i++) {
                long sessionId = ia.readLong("sessionId");
                TreeMap<Integer, ByteBuffer> client = new TreeMap<Integer, ByteBuffer>();
                int clientSize = ia.readInt("clientsize");
                System.out.println("deserialize client " + sessionId + " " + clientSize);
                for (int j = 0; j < clientSize; j++) {
                    int xid = ia.readInt("xid");
                    ByteBuffer bb = ByteBuffer.wrap(ia.readBuffer("data"));
                    client.put(xid, bb);
		    if(Configuration.debug)
                        System.out.println("deserialize " + sessionId + " " + xid);
                }

                cache.put(sessionId, client);
            }
        //}
    }

    private static HashMap<Long, Integer> lastOrdered = new HashMap<Long, Integer>();

    public static int getLastOrdered(long sessionId){
        synchronized(lastOrdered){
            if(!lastOrdered.containsKey(sessionId))
                return -1;
            else
                return lastOrdered.get(sessionId);
        }
    }

    public static void updateLastOrdered(long sessionId, int cxid){
        synchronized(lastOrdered){
            lastOrdered.put(sessionId, cxid);
        }
    }

    public synchronized static void rebuild(){
        lastOrdered.clear();
        for(Long sessionId : cache.keySet()){
            lastOrdered.put(sessionId, cache.get(sessionId).lastKey());    
        }
    }

}
